"""
分布式追踪模块 - 基于OpenTelemetry

这个模块实现了分布式追踪功能：
- 为每个任务生成Trace ID
- 跨服务链路追踪
- 性能监控和延迟分析
- 集成Jaeger可视化
"""

import asyncio
import logging
import time
import uuid
from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Union
from datetime import datetime, timezone

# OpenTelemetry imports
from opentelemetry import trace
from opentelemetry.exporter.jaeger.thrift import JaegerExporter
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.sdk.resources import Resource
from opentelemetry.trace import Status, StatusCode
from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
from opentelemetry.instrumentation.requests import RequestsInstrumentor
from opentelemetry.instrumentation.aiohttp_client import AioHttpClientInstrumentor
from opentelemetry.instrumentation.flask import FlaskInstrumentor
from opentelemetry.instrumentation.asyncio import AsyncioInstrumentor

logger = logging.getLogger(__name__)


@dataclass
class TraceConfig:
    """追踪配置"""
    service_name: str = "gpu-auto-scaling-system"
    service_version: str = "2.0.0"
    jaeger_endpoint: str = "http://localhost:14268/api/traces"
    environment: str = "development"
    enable_auto_instrumentation: bool = True
    enable_jaeger: bool = True
    sample_rate: float = 1.0  # 采样率，1.0表示100%采样
    max_export_batch_size: int = 512
    export_timeout_millis: int = 30000
    schedule_delay_millis: int = 5000


class DistributedTracer:
    """分布式追踪器"""
    
    def __init__(self, config: Optional[TraceConfig] = None):
        self.config = config or TraceConfig()
        self.tracer_provider: Optional[TracerProvider] = None
        self.tracer = None
        self.propagator = TraceContextTextMapPropagator()
        self._initialized = False
        
        logger.info("分布式追踪器已初始化")
    
    def initialize(self) -> None:
        """初始化追踪器"""
        if self._initialized:
            return
        
        try:
            # 创建资源
            resource = Resource.create({
                "service.name": self.config.service_name,
                "service.version": self.config.service_version,
                "deployment.environment": self.config.environment,
            })
            
            # 创建追踪器提供者
            self.tracer_provider = TracerProvider(
                resource=resource,
                sampler=trace.sampling.TraceIdRatioBased(self.config.sample_rate)
            )
            
            # 设置全局追踪器提供者
            trace.set_tracer_provider(self.tracer_provider)
            
            # 创建追踪器
            self.tracer = trace.get_tracer(
                self.config.service_name,
                self.config.service_version
            )
            
            # 配置Jaeger导出器
            if self.config.enable_jaeger:
                jaeger_exporter = JaegerExporter(
                    agent_host_name="localhost",
                    agent_port=6831,
                    collector_endpoint=self.config.jaeger_endpoint,
                )
                
                span_processor = BatchSpanProcessor(
                    jaeger_exporter,
                    max_export_batch_size=self.config.max_export_batch_size,
                    export_timeout_millis=self.config.export_timeout_millis,
                    schedule_delay_millis=self.config.schedule_delay_millis,
                )
                
                self.tracer_provider.add_span_processor(span_processor)
            
            # 启用自动插桩
            if self.config.enable_auto_instrumentation:
                self._enable_auto_instrumentation()
            
            self._initialized = True
            logger.info("分布式追踪器初始化成功")
            
        except Exception as e:
            logger.error(f"分布式追踪器初始化失败: {e}")
            raise
    
    def _enable_auto_instrumentation(self) -> None:
        """启用自动插桩"""
        try:
            # HTTP客户端插桩
            RequestsInstrumentor().instrument()
            AioHttpClientInstrumentor().instrument()
            
            # 异步操作插桩
            AsyncioInstrumentor().instrument()
            
            logger.info("自动插桩已启用")
            
        except Exception as e:
            logger.warning(f"启用自动插桩失败: {e}")
    
    def create_trace_context(self, 
                           task_id: str,
                           operation_name: str,
                           parent_context: Optional[Dict[str, str]] = None) -> Dict[str, str]:
        """创建追踪上下文"""
        if not self._initialized:
            self.initialize()
        
        # 创建新的span
        with self.tracer.start_as_current_span(operation_name) as span:
            # 设置span属性
            span.set_attribute("task.id", task_id)
            span.set_attribute("operation.name", operation_name)
            span.set_attribute("timestamp", datetime.now(timezone.utc).isoformat())
            
            # 创建传播上下文
            context = {}
            self.propagator.inject(context)
            
            return context
    
    def extract_trace_context(self, context: Dict[str, str]) -> Optional[Any]:
        """提取追踪上下文"""
        if not self._initialized:
            return None
        
        try:
            return self.propagator.extract(context)
        except Exception as e:
            logger.warning(f"提取追踪上下文失败: {e}")
            return None
    
    @contextmanager
    def start_span(self, 
                   operation_name: str,
                   task_id: Optional[str] = None,
                   attributes: Optional[Dict[str, Any]] = None,
                   parent_context: Optional[Dict[str, str]] = None):
        """开始一个span"""
        if not self._initialized:
            self.initialize()
        
        # 提取父上下文
        parent_span_context = None
        if parent_context:
            parent_span_context = self.extract_trace_context(parent_context)
        
        # 创建span
        span = self.tracer.start_span(
            operation_name,
            context=parent_span_context
        )
        
        try:
            # 设置属性
            if task_id:
                span.set_attribute("task.id", task_id)
            
            if attributes:
                for key, value in attributes.items():
                    span.set_attribute(key, value)
            
            yield span
            
        except Exception as e:
            # 标记span为错误状态
            span.set_status(Status(StatusCode.ERROR, str(e)))
            span.set_attribute("error", True)
            span.set_attribute("error.message", str(e))
            raise
        finally:
            span.end()
    
    @asynccontextmanager
    async def start_async_span(self, 
                              operation_name: str,
                              task_id: Optional[str] = None,
                              attributes: Optional[Dict[str, Any]] = None,
                              parent_context: Optional[Dict[str, str]] = None):
        """开始一个异步span"""
        if not self._initialized:
            self.initialize()
        
        # 提取父上下文
        parent_span_context = None
        if parent_context:
            parent_span_context = self.extract_trace_context(parent_context)
        
        # 创建span
        span = self.tracer.start_span(
            operation_name,
            context=parent_span_context
        )
        
        try:
            # 设置属性
            if task_id:
                span.set_attribute("task.id", task_id)
            
            if attributes:
                for key, value in attributes.items():
                    span.set_attribute(key, value)
            
            yield span
            
        except Exception as e:
            # 标记span为错误状态
            span.set_status(Status(StatusCode.ERROR, str(e)))
            span.set_attribute("error", True)
            span.set_attribute("error.message", str(e))
            raise
        finally:
            span.end()
    
    def add_event(self, 
                  span: Any,
                  event_name: str,
                  attributes: Optional[Dict[str, Any]] = None) -> None:
        """添加事件到span"""
        if not span:
            return
        
        try:
            span.add_event(event_name, attributes or {})
        except Exception as e:
            logger.warning(f"添加事件失败: {e}")
    
    def set_span_status(self, 
                       span: Any,
                       status_code: StatusCode,
                       description: Optional[str] = None) -> None:
        """设置span状态"""
        if not span:
            return
        
        try:
            span.set_status(Status(status_code, description))
        except Exception as e:
            logger.warning(f"设置span状态失败: {e}")
    
    def get_current_span(self) -> Optional[Any]:
        """获取当前span"""
        if not self._initialized:
            return None
        
        try:
            return trace.get_current_span()
        except Exception:
            return None
    
    def get_trace_id(self) -> Optional[str]:
        """获取当前追踪ID"""
        span = self.get_current_span()
        if span and span.get_span_context().is_valid:
            return format(span.get_span_context().trace_id, '032x')
        return None
    
    def get_span_id(self) -> Optional[str]:
        """获取当前span ID"""
        span = self.get_current_span()
        if span and span.get_span_context().is_valid:
            return format(span.get_span_context().span_id, '016x')
        return None
    
    def shutdown(self) -> None:
        """关闭追踪器"""
        if self.tracer_provider:
            self.tracer_provider.shutdown()
            logger.info("分布式追踪器已关闭")


class TaskTracer:
    """任务追踪器 - 专门用于任务追踪"""
    
    def __init__(self, distributed_tracer: DistributedTracer):
        self.distributed_tracer = distributed_tracer
        self.logger = logging.getLogger(__name__)
    
    def create_task_trace(self, 
                         task_id: str,
                         operation: str,
                         attributes: Optional[Dict[str, Any]] = None) -> Dict[str, str]:
        """创建任务追踪"""
        trace_attributes = {
            "task.id": task_id,
            "task.operation": operation,
            "task.start_time": datetime.now(timezone.utc).isoformat()
        }
        
        if attributes:
            trace_attributes.update(attributes)
        
        return self.distributed_tracer.create_trace_context(
            task_id=task_id,
            operation_name=f"task.{operation}",
            parent_context=None
        )
    
    @contextmanager
    def trace_task_operation(self, 
                            task_id: str,
                            operation: str,
                            attributes: Optional[Dict[str, Any]] = None):
        """追踪任务操作"""
        trace_attributes = {
            "task.id": task_id,
            "task.operation": operation
        }
        
        if attributes:
            trace_attributes.update(attributes)
        
        with self.distributed_tracer.start_span(
            operation_name=f"task.{operation}",
            task_id=task_id,
            attributes=trace_attributes
        ) as span:
            yield span
    
    @asynccontextmanager
    async def trace_async_task_operation(self, 
                                        task_id: str,
                                        operation: str,
                                        attributes: Optional[Dict[str, Any]] = None):
        """追踪异步任务操作"""
        trace_attributes = {
            "task.id": task_id,
            "task.operation": operation
        }
        
        if attributes:
            trace_attributes.update(attributes)
        
        async with self.distributed_tracer.start_async_span(
            operation_name=f"task.{operation}",
            task_id=task_id,
            attributes=trace_attributes
        ) as span:
            yield span
    
    def add_task_event(self, 
                      span: Any,
                      event_name: str,
                      task_id: str,
                      attributes: Optional[Dict[str, Any]] = None) -> None:
        """添加任务事件"""
        event_attributes = {
            "task.id": task_id,
            "event.name": event_name
        }
        
        if attributes:
            event_attributes.update(attributes)
        
        self.distributed_tracer.add_event(span, event_name, event_attributes)
    
    def mark_task_success(self, 
                         span: Any,
                         task_id: str,
                         result_url: Optional[str] = None,
                         processing_duration: Optional[float] = None) -> None:
        """标记任务成功"""
        if span:
            span.set_attribute("task.status", "success")
            if result_url:
                span.set_attribute("task.result_url", result_url)
            if processing_duration:
                span.set_attribute("task.processing_duration", processing_duration)
            
            self.distributed_tracer.set_span_status(span, StatusCode.OK)
    
    def mark_task_failure(self, 
                         span: Any,
                         task_id: str,
                         error_message: str,
                         processing_duration: Optional[float] = None) -> None:
        """标记任务失败"""
        if span:
            span.set_attribute("task.status", "failure")
            span.set_attribute("task.error_message", error_message)
            if processing_duration:
                span.set_attribute("task.processing_duration", processing_duration)
            
            self.distributed_tracer.set_span_status(span, StatusCode.ERROR, error_message)


# 全局追踪器实例
_global_tracer: Optional[DistributedTracer] = None
_global_task_tracer: Optional[TaskTracer] = None


def get_tracer() -> DistributedTracer:
    """获取全局追踪器实例"""
    global _global_tracer
    if _global_tracer is None:
        _global_tracer = DistributedTracer()
    return _global_tracer


def get_task_tracer() -> TaskTracer:
    """获取全局任务追踪器实例"""
    global _global_task_tracer
    if _global_task_tracer is None:
        _global_task_tracer = TaskTracer(get_tracer())
    return _global_task_tracer


def initialize_tracing(config: Optional[TraceConfig] = None) -> None:
    """初始化全局追踪"""
    global _global_tracer, _global_task_tracer
    
    _global_tracer = DistributedTracer(config)
    _global_tracer.initialize()
    
    _global_task_tracer = TaskTracer(_global_tracer)
    
    logger.info("全局追踪已初始化")


def shutdown_tracing() -> None:
    """关闭全局追踪"""
    global _global_tracer
    if _global_tracer:
        _global_tracer.shutdown()
        _global_tracer = None
